#！/bin/bash

current_dir=$(cd $(dirname $0); pwd)
echo $current_dir
workdir=$current_dir

cd $workdir


HOST=${HOST:=0.0.0.0}
PORT=${PORT:=8080}
WORLD_SIZE=${WORLD_SIZE:=2}
PIPEFUSION_PARALLEL_DEGREE=${PIPEFUSION_PARALLEL_DEGREE:=1}
ULYSSES_DEGREE=${ULYSSES_DEGREE:=1}
RING_DEGREE=${RING_DEGREE:=1}
USE_CFG_PARALLEL=${USE_CFG_PARALLEL:=1}
USE_TORCH_COMPILE=${USE_TORCH_COMPILE:=0}
IF_LLM=${IF_LLM_PROMPT:=0}
IF_SUPER=${IF_SUPER_RESOLUTION:=0}
SUPER_UPSCALE=${SUPER_UPSCALE:=2}
FACE_UPSAMPLE=${FACE_UPSAMPLE:=0}
USE_TAYLORSEER=${USE_TAYLORSEER:=1}
USE_QUANTIZATION_INT8=${USE_QUANTIZATION_INT8:=0}
MAX_CONCURRENCY_SIZE=${MAX_CONCURRENCY_SIZE:=1}
MAX_PROCESS_TASK_SIZE=${MAX_PROCESS_TASK_SIZE:=3}

if [ -z ${MODEL_PATH} ];then
    echo "environment var MODEL_PATH not set"
    exit 1
fi


cmd="python serving.py \
--host ${HOST} \
--port ${PORT} \
--model ${MODEL_PATH} \
--max_concurrency_size ${MAX_CONCURRENCY_SIZE} \
--max_process_task_size ${MAX_PROCESS_TASK_SIZE} \
--world_size ${WORLD_SIZE} \
--pipefusion_parallel_degree ${PIPEFUSION_PARALLEL_DEGREE} \
--ulysses_degree ${ULYSSES_DEGREE} \
--ring_degree ${RING_DEGREE} \
--upscale ${SUPER_UPSCALE} \
--if_llm ${IF_LLM} \
--face_upsample ${FACE_UPSAMPLE} \
--api_key ${API_KEY} \
--base_url ${BASE_URL} \
--llm_model ${LLM_MODEL} \
--if_super ${IF_SUPER} \
--upscale ${SUPER_UPSCALE} \
--use_taylorseer ${USE_TAYLORSEER} \
--use_quantization_int8 ${USE_QUANTIZATION_INT8} \
--use_cfg_parallel"


# if [ "$USE_CFG_PARALLEL" -eq 1 ]; then
#     cmd+=" --use_cfg_parallel"
# fi
if [ "$USE_TORCH_COMPILE" -eq 1 ]; then
    # cmd+=" --use_torch_compile"
    cmd="python serving.py \
--host ${HOST} \
--port ${PORT} \
--max_concurrency_size ${MAX_CONCURRENCY_SIZE} \
--max_process_task_size ${MAX_PROCESS_TASK_SIZE} \
--model ${MODEL_PATH} \
--world_size ${WORLD_SIZE} \
--pipefusion_parallel_degree ${PIPEFUSION_PARALLEL_DEGREE} \
--ulysses_degree ${ULYSSES_DEGREE} \
--ring_degree ${RING_DEGREE} \
--upscale ${SUPER_UPSCALE} \
--if_llm ${IF_LLM} \
--face_upsample ${FACE_UPSAMPLE} \
--api_key ${API_KEY} \
--base_url ${BASE_URL} \
--llm_model ${LLM_MODEL} \
--if_super ${IF_SUPER} \
--upscale ${SUPER_UPSCALE} \
--use_cfg_parallel \
--use_taylorseer ${USE_TAYLORSEER} \
--use_quantization_int8 ${USE_QUANTIZATION_INT8} \
--use_torch_compile"
fi


# if [ "$USE_CFG_PARALLEL" -eq 1 ]; then
#     cmd+=" --use_cfg_parallel"
# fi


if [ "$IF_LLM" -eq 1 ]; then

    if [ -z ${API_KEY} ];then
        echo "environment var API_KEY not set"
        exit 1
    fi

    if [ -z ${BASE_URL} ];then
        echo "environment var BASE_URL not set"
        exit 1
    fi

    if [ -z ${LLM_MODEL} ];then
        echo "environment var LLM_MODEL not set"
        exit 1
    fi
fi

if [ "$IF_SUPER" -eq 1 ]; then
    cmd+=" --if_super"
fi

if [ "$FACE_UPSAMPLE" -eq 1 ]; then
    cmd+=" --face_upsample"
fi

echo $cmd

$cmd